import random

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.preprocessing import StandardScaler
import pandas as pd

# Load data
df = pd.read_csv("static/sensor_data.csv")

# Prepare training and validation data
scaler_x = StandardScaler()
scaler_y = StandardScaler()

X_train = scaler_x.fit_transform(df.iloc[:100, 1:-2].values)
y_train = scaler_y.fit_transform(df.iloc[:100, -2:].values)

X_val = scaler_x.transform(df.iloc[100:200, 1:-2].values)
y_val = scaler_y.transform(df.iloc[100:200, -2:].values)

X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)
X_val = torch.tensor(X_val, dtype=torch.float32)
y_val = torch.tensor(y_val, dtype=torch.float32)

# Define neural network
class SimpleNN(nn.Module):
    def __init__(self, input_size, output_size):
        super(SimpleNN, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(input_size, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, output_size)
        )

    def forward(self, x):
        return self.net(x)

if __name__ == '__main__':
    torch.manual_seed(42)
    torch.cuda.manual_seed_all(42)
    model = SimpleNN(input_size=13, output_size=2)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.0005)

    # Training loop
    for epoch in range(3000):
        model.train()
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()

        if (epoch + 1) % 20 == 0:
            model.eval()
            val_outputs = model(X_val)
            val_loss = criterion(val_outputs, y_val)
            print(f"Epoch {epoch+1}, Training Loss: {loss.item():.4f}, Validation Loss: {val_loss.item():.4f}")

    # Save model
    torch.save(model.state_dict(), "model.pth")


    # Evaluate and plot predictions vs ground truth
    import matplotlib.pyplot as plt

    model.eval()
    predictions = model(X_val).detach().numpy()
    ground_truth = y_val.numpy()

    # Inverse transform to original scale
    predictions = scaler_y.inverse_transform(predictions)
    ground_truth = scaler_y.inverse_transform(ground_truth)

    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.plot(ground_truth[:, 0], label='True FIC1008')
    plt.plot(predictions[:, 0], label='Predicted FIC1008')
    plt.title('FIC1008 Prediction')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(ground_truth[:, 1], label='True FI615')
    plt.plot(predictions[:, 1], label='Predicted FI615')
    plt.title('FI615 Prediction')
    plt.legend()

    plt.tight_layout()
    plt.show()
